from ast import While
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import gc
class Point:

    def __init__(self,point):
        self.parentIndex=-1
        self.children=[]
        self.parent=None
        self.index=int(point[0])
        self.type=int(point[1])
        self.x=point[2]
        self.y=point[3]
        self.z=point[4]
        self.ratio=point[5]
        self.parentIndex=int(point[6])
        self.children=[]
        self.xyz=[np.float32(point[2]),np.float32(point[3]),np.float32(point[4])]
    def __str__(self):
        
        return str(self.index)+","+str(self.type)+","+str(self.x)+","+\
        str(self.y)+","+str(self.z)+","+str(self.ratio)+","+str(self.parentIndex)
        
class Edge:
    def __init__(self,point=None):
        if point is None:
            self.data=[]
        else:
            self.data=[point]
        self.order=-1
        self.angle=0
        self.maxDepth=-1
        self.maxLength=-1
        self.children=[]
        self.parent=None
        self.len=0
    def addPoint(self,point):
        self.data.append(point)
    def getLength(self):
        if self.len==0:
            for i in range(len(self.data)-1):
                self.len+=np.linalg.norm(np.array(self.data[i].xyz)-np.array(self.data[i+1].xyz))
        return self.len
        
class NeuronTree:
    def __init__(self):
        self.root=None
        self.edges=[]
        self.points=[]
        self.width=1
        self.branchs=[]
        self.terminals=[]
        self.swc={}
        self.xyz=[]
        self.somaRadius=200
        self.somaHide=False
        self.axonHide=False
        self.dendriteHide=False
        self.maxOrder=-1
        self.rootAxonEdge=None
        self.dendrites=[]
 
    def readSWC(self,swc):
        swclines=swc.split('\n')
        for line in swclines:
            if len(line)==0 or line[0] == '#'or line[0] == '\r':
                continue
            p=list(map(float,(line.split())))
            self.swc[int(p[0])]=Point(p)
        # self.swc.pop()
        self.parseswc()

    def readFile(self,filename='',spliter=' '):

        with open(filename, "r") as f:
            text_lines = f.read()
            swclines=text_lines.split('\n')
            for line in swclines:
                if len(line)==0 or line[0] == '#':
                    continue
                p=list(map(float,(line.split(spliter))))
                self.swc[int(p[0])]=Point(p)
        self.parseswc()

    def parseswc(self):
        for key,point in self.swc.items():
            self.xyz.append(point.xyz)
            self.points.append(point)
            if point.parentIndex==-1:
                if self.root is None:
                    self.root=point
        self.getEdges()
        for edge in self.edges:
            if self.rootAxonEdge is None and (edge.data[0]==self.root or edge.data[0].type==3)and (edge.data[1].type!=3):
                self.rootAxonEdge=edge
            if len(edge.data[len(edge.data)-1].children)==0:
                self.terminals.append(edge.data[len(edge.data)-1])
        for edge in self.edges:
            for tail in edge.data[len(edge.data)-1].children:
                for edge0 in self.edges:
                    if tail==edge0.data[1]:
                        edge.children.append(edge0)
                        edge0.parent = edge
    def getDendrite(self):
        if len(self.dendrites)==0:
            for edge in self.edges:
                if edge.data[1].type!=3:
                    continue
                self.dendrites.append(edge)
        return self.dendrites
        
    def getEdges(self):
        for index,point in self.swc.items():
            if point.parentIndex!=-1  and self.swc.get(point.parentIndex) is not None:
                point.parent=self.swc[int(point.parentIndex)]
                point.parent.children.append(point)
            pass
        start = self.swc[list(self.swc.keys())[0]]
        if len(start.children):
            lastp=None
            self.getEdge(start)

    def getEdge(self,start):
        for p in start.children:
            edge=Edge(start)
            edgestart=p
            while len(edgestart.children)==1:
                edge.addPoint(edgestart)
                #del self.swc[edgestart.index]
                edgestart=edgestart.children[0]
            edge.addPoint(edgestart)
            self.edges.append(edge)
            lastp=edgestart
            self.getEdge(lastp)
        pass
    def getEdgeByTerminal(self,terminal):
        for edge in self.edges:
            if edge.data[len(edge.data)-1]==terminal:
                return edge
        pass
    def print(self):
        level='-';
        print(level,self.root.index)
        self.printChildren(self.root,level)
        
        
    def printChildren(self,point,level):
        level=level+'---'
        for child in point.children:
            print(level,child.index)
            self.printChildren(child,level)
            
    def printEdges(self):
        for edge in self.edges:
            print('==============')
            for p in edge:
                print (p.index);

    def plotMPR(self,filename=None):
        dpi = 100
        fig = plt.figure(figsize=(1200/dpi, 400/dpi), dpi=dpi)
        
        axes0 = plt.subplot(131)
        plt.axis('off')
        axes0.set_xlim(0, 11400)
        axes0.set_ylim(0, 10000)
        fig.patch.set_facecolor('#000000')
        # plt.Circle((neuron.root.z, neuron.root.y), 300)
        mirror = False
        if self.root.z<5700:
            mirror=True
        for edge in self.edges:
            x=[]
            y=[]
            for p in edge.data:
                z= p.z
                if mirror:
                    z= 11400-p.z
                x.append(11400-z)
                y.append(10000-p.y)
            plt.plot(x, y,'g-',color='#FFFFFF')
        # plt.figure(facecolor='gainsboro')
        rootz = self.root.z
        if mirror:
            rootz= 11400-self.root.z
        plt.plot(10000-rootz,10000- self.root.y,'g,')

        axes0 =  plt.subplot(132)
        axes0.set_axis_off()
        axes0.set_xlim(0, 16200)
        axes0.set_ylim(0, 16200)
        fig.patch.set_facecolor('#000000')
        # plt.Circle((neuron.root.z, neuron.root.y), 300)
        for edge in self.edges:
            x=[]
            y=[]
            for p in edge.data:
                x.append(16200-p.x)
                y.append(10000-p.y)
            plt.plot(x, y,'g-',color='#FFFFFF')
        # plt.figure(facecolor='gainsboro')
        plt.plot(10000-self.root.x,10000- self.root.y,'g,')

        axes0 = plt.subplot(133)
        plt.axis('off')
        axes0.set_xlim(0, 16200)
        axes0.set_ylim(0, 16200)
        fig.patch.set_facecolor('#000000')
        # plt.Circle((neuron.root.z, neuron.root.y), 300)
        for edge in self.edges:
            x=[]
            y=[]
            for p in edge.data:
                z= p.z
                if mirror:
                    z= 11400-p.z
                x.append(16200-p.x)
                y.append(11400-z)
            plt.plot(x, y,'g-',color='#FFFFFF')
        # plt.figure(facecolor='gainsboro')
        plt.plot(10000-self.root.x,10000- rootz,'g,')

        fname = filename
        if filename is None:
            fname="../resource/neuron.png"
        plt.savefig(fname=fname,format="png",bbox_inches='tight')
        fig.clf()
        plt.clf()
        plt.cla()
        plt.close('all')
        plt.close('fig')
        gc.collect()
        pass
        
if __name__ == "__main__":
    import os,sys
    if hasattr(sys.modules[__name__], '__file__'):
        _file_name = __file__
    else:
        _file_name = inspect.getfile(inspect.currentframe())
    CURRENT_FILE_PATH = os.path.dirname(_file_name)

    tree = NeuronTree()
    # tree.readFile("../resource/swc/17109/17109_1701_x8048_y22277.semi_r.swc")
    tree.readFile(CURRENT_FILE_PATH+"/../resource/swc/18715/025.swc")
    tree.plotMPR()
    print(len(tree.edges))
    # for edge in tree.edges:
    #     print(len(edge.data))
        # for p in edge.data:
        #     print(p)
